#include "star_stack.h"

#include "dwarf_utils.h"
#include "elog.h"
#include "elog_file.h"

#include <boost/math/interpolators/pchip.hpp>

using namespace cv;

void rotateImage(Mat src, Mat& dst, float angle) {
	if(src.empty())
		return;

	int h = src.rows;//定义图片的高度
	int w = src.cols;//定义图片的宽度

	Point2f center((float)w / 2.0, (float)h / 2.0);

	Mat mapMatrix = getRotationMatrix2D(center, angle, 1.0);
	double cos = abs(mapMatrix.at<double>(0, 0));
	double sin = abs(mapMatrix.at<double>(0, 1));
	int nw = cos * w + sin * h;
	int nh = sin * w + cos * h;
	mapMatrix.at<double>(0, 2) += (nw - w) / 2.0;
	mapMatrix.at<double>(1, 2) += (nh - h) / 2.0;
	warpAffine(src, dst, mapMatrix, Size(nw, nh));
}

void scaleImage(const Mat& src, Mat& dst, float scale) {
	if(src.empty())
		return;
		
	Point2f center(0.0, 0.0);

	Mat mapMatrix = getRotationMatrix2D(center, 0.0, scale);
	warpAffine(src, dst, mapMatrix, dst.size());
}

// 耗时和旧的一样，但写法简单很多
int stackImageNew(const Mat& img, Mat& imgStack, Mat& imgStackAll, Mat& normalCounter, const starMatch* starMatches) {	 
	if(img.empty())
		return -1;

	if(imgStackAll.empty()) 
		return -1;

	if(normalCounter.empty()) 
		return -1;

	Point2f center((float)img.cols / 2.0, (float)img.rows / 2.0);
	Mat mapMatrix = getRotationMatrix2D(center, starMatches[0].rotate * (180.0 / M_PI), 1.0);// 先旋转
	mapMatrix.at<double>(0, 2) += starMatches[0].deltaX;// 再平移回中心
	mapMatrix.at<double>(1, 2) += starMatches[0].deltaY;

	Mat imgTemp;
	warpAffine(img, imgTemp, mapMatrix, img.size());
	
	imgStack = cv::Mat::zeros(img.rows, img.cols, CV_16UC3);

	float* img_stack_all_ptr = reinterpret_cast<float*>(imgStackAll.data);
	uint16_t* img_stack_ptr = reinterpret_cast<uint16_t*>(imgStack.data);
	uint16_t* normal_counter_ptr = reinterpret_cast<uint16_t*>(normalCounter.data);
	uint16_t* img_temp_ptr = reinterpret_cast<uint16_t*>(imgTemp.data);

	size_t data_size = imgStackAll.cols * imgStackAll.rows * imgStackAll.channels();
	for (size_t i = 0; i < data_size; ++i) {
		img_stack_all_ptr[i] += img_temp_ptr[i];
		normal_counter_ptr[i] += img_temp_ptr[i] > 0 ? 1 : 0;
		img_stack_ptr[i] = img_stack_all_ptr[i] / normal_counter_ptr[i];
	}

	return 0;
} 

int stackImage(const Mat& img, Mat& imgStack, Mat& imgStackAll, Mat& normalCounter, const starMatch* starMatches) {	 
	if(img.empty())
		return -1;

	if(imgStackAll.empty()) 
		return -1;

	if(normalCounter.empty()) 
		return -1;

	Mat imgTemp;
	rotateImage(img, imgTemp, starMatches[0].rotate * (180.0 / M_PI));

	Point2i stackOffset;
	stackOffset.x = (imgTemp.cols - img.cols) / 2 - starMatches[0].deltaX;
	stackOffset.y = (imgTemp.rows - img.rows) / 2 - starMatches[0].deltaY;

	Point2i imgStackAllOffset;
	imgStackAllOffset.x = stackOffset.x < 0 ? 0 : stackOffset.x;
	imgStackAllOffset.y = stackOffset.y < 0 ? 0 : stackOffset.y;

	Point2i imgTempOffset;
	imgTempOffset.x = imgStackAllOffset.x - stackOffset.x;
	imgTempOffset.y = imgStackAllOffset.y - stackOffset.y;

	Rect imgTempRect(imgTempOffset, Size(imgTemp.cols, imgTemp.rows));
	Rect imgStackAllRect(imgStackAllOffset, Size(img.cols, img.rows));
	Rect intersect = imgTempRect & imgStackAllRect;
	if(intersect.area() > 0) {
		Point2i intersectInimgTempOffset;
		intersectInimgTempOffset.x = intersect.tl().x - imgTempRect.tl().x;
		intersectInimgTempOffset.y = intersect.tl().y - imgTempRect.tl().y;
		Rect intersectInimgTemp(intersectInimgTempOffset, intersect.size());

		Point2i intersectInimgStackAllOffset;
		intersectInimgStackAllOffset.x = intersect.tl().x - imgStackAllRect.tl().x;
		intersectInimgStackAllOffset.y = intersect.tl().y - imgStackAllRect.tl().y;
		Rect intersectInimgStackAll(intersectInimgStackAllOffset, intersect.size());

		//add the image here
		imgStackAll(intersectInimgStackAll) += imgTemp(intersectInimgTemp);

		cv::threshold(imgTemp, imgTemp, 0, 1, THRESH_BINARY);
		normalCounter(intersectInimgStackAll) += imgTemp(intersectInimgTemp);

		imgStack = imgStackAll / normalCounter;
		imgStack.convertTo(imgStack, CV_16UC3);

		return 0;
	} else {
		return -2;
	}
} 

int getMidValue(const Mat& input)
{
	if(input.empty())
		return -1;

	int rows = input.rows;
	int cols = input.cols;
	float histogram[65536] = {0};
    //先计算图像的直方图
	for (int i = 0; i < rows; ++i) {
		///获取i行首像素的指针
		const ushort *p = input.ptr<ushort>(i);
		///遍历i行像素
		for (int j = 0; j < cols; ++j) {
			//cout << int(*p++) << endl;
			histogram[int(*p++)]++;
		}
	}
	int HalfNum = rows * cols / 2;
	int tempSum = 0;
	for (int i = 0; i < 65536; i++) {
		tempSum = tempSum + histogram[i];
		if (tempSum > HalfNum)
			return i;
	}
	return -1;
}

double midtonesTransferFunction(double midToneBalance, double x) 
{
	if (x > 0) {
		if (x < 1) {
			return (midToneBalance - 1) * x / ((2 * midToneBalance - 1) * x - midToneBalance);
		}
		return 1;
	}
	return 0;
}

double normalizeUShort(double val, double maxVal) {
	return val / maxVal;
}

ushort denormalizeUShort(double val, double maxVal) {
	return (ushort)(val * maxVal);
}

int getStretchMap(Mat &img, double targetHistogramMedianPercent, double shadowsClipping) {
	if(img.empty())
		return -1;
	double maxVal = 0;
	double minVal = 0;
	minMaxLoc(img, &minVal, &maxVal);
	// log_i("maxVal:%f", maxVal);

	int median = getMidValue(img);
	// log_i("median:%d", median);
	double normalizedMedian = normalizeUShort(median, maxVal);
	// log_i("normalizedMedian:%f", normalizedMedian);

	Mat imgTemp = cv::abs(img - cv::Scalar(median,median));
	int medianAbsoluteDeviation = getMidValue(imgTemp);
	// log_i("mad:%d", medianAbsoluteDeviation);
	double normalizedMAD = normalizeUShort(medianAbsoluteDeviation, maxVal);
	// log_i("normalizedMAD:%f", normalizedMAD);

	double scaleFactor = 1.4826; // see https://en.wikipedia.org/wiki/Median_absolute_deviation

	double shadows = 0.0;
	double midtones = 0.5;
	double highlights = 1.0;

	shadows = normalizedMedian + shadowsClipping * normalizedMAD * scaleFactor;
	midtones = midtonesTransferFunction(targetHistogramMedianPercent, normalizedMedian - shadows);
	highlights = 1;
	// log_i("shadows:%f", shadows);
	// log_i("midtones:%f", midtones);
	// log_i("highlights:%f", highlights);

	img.forEach<ushort>([&](ushort& pixel, const int position[]) {
		double value = normalizeUShort(pixel, maxVal);
		pixel = denormalizeUShort(midtonesTransferFunction(midtones, 1 - highlights + value - shadows), maxVal);
	});

	return 0;
}

int unlinkedStretch(Mat &img) {
	if(img.empty())
		return -1;
	if(img.channels() != 3)
		return -2;

	Mat mv[3];
	double autoStretchFactor = 0.2;
	double blackClipping = -2.8;
	
	split(img, mv);
	for(int k = 0; k < 3; k++) 
		getStretchMap(mv[k], autoStretchFactor, blackClipping);
	merge(mv, 3, img);

	return 0;
}

double median(cv::Mat channel, int nvals, float upper_range) {
    double m = (channel.rows*channel.cols) / 2;
    int bin = 0;
    double med = -1.0;

    int histSize = nvals;
    float range[] = {0, upper_range};
    const float* histRange = { range };
    bool uniform = true;
    bool accumulate = false;
    cv::Mat hist;
    cv::calcHist(&channel, 1, 0, cv::Mat(), hist, 1, &histSize, &histRange, uniform, accumulate);

    for (int i = 0; i < histSize && med < 0.0; ++i) {
        bin += cvRound(hist.at<float>(i));
        if (bin > m && med < 0.0)
            med = i;
    }

    return med;
}

int imgSegSStretch(const cv::Mat& img_dmax_norm, cv::Mat& img_binary) {
    cv::Mat img_uint8;
    img_dmax_norm.convertTo(img_uint8, CV_8U, 255.0);

    cv::Mat img_HSV;
    cv::cvtColor(img_uint8, img_HSV, cv::COLOR_BGR2HSV);

    cv::Mat img_V;
    cv::extractChannel(img_HSV, img_V, 2);

    cv::threshold(img_V, img_binary, 0, 1, cv::THRESH_TRIANGLE);

    int count_non_zero = cv::countNonZero(img_binary);
    double img_binary_area_ratio = count_non_zero / (double)img_V.total();

    if (img_binary_area_ratio < 0.2) {
        img_binary.convertTo(img_binary, CV_32F);
    } else {
        cv::Mat signal_pixels = img_V.mul(img_binary);

        cv::Mat signal_pixels_non_zero(cv::Size(1, count_non_zero), CV_8U);
        uchar* signal_pixels_ptr_uchar = signal_pixels_non_zero.ptr<uchar>(0);

        std::mutex non_zero_mutex;

        int idx = 0;
        signal_pixels.forEach<uchar>([&](uchar& pixel, const int position[]) {
            if (pixel != 0) {
                std::lock_guard<std::mutex> lock(non_zero_mutex);
                signal_pixels_ptr_uchar[idx] = pixel;
                ++idx;
            }
        });

        double temp_median = median(signal_pixels_non_zero, 256, 256);
        cv::Scalar temp_stddev;
        cv::meanStdDev(signal_pixels_non_zero, cv::noArray(), temp_stddev);

        double th = temp_median + 2.0 * temp_stddev.val[0];

        std::cout << "th: " << th << std::endl;
        std::cout << "temp_median: " << temp_median << std::endl;
        std::cout << "temp_stddev.val[0]: " << temp_stddev.val[0] << std::endl;

        cv::threshold(img_V, img_binary, th, 1, cv::THRESH_BINARY);
        img_binary.convertTo(img_binary, CV_32F);
    }

    return 0;
}

int imgSCurveHistogramTransformUseV(const cv::Mat& img_float32, const cv::Mat& img_BGR_norm,
                                          const cv::Mat& star_mask, cv::Mat& img_transfer) {
    const double point1_LR = 0.05;   // 控制阴影保护点的x坐标 0.05
    const double point1_UD = 0.40;   // 控制阴影影保护点的y坐标 0.45
    const double point2_LR = 0.25;   // 控制亮部保护点的x坐标，该值设置的小一些可以暴露淡运气但是背景不均匀部分颜色可能会不对 0.8
    const double point2_UP = 0.05;   // 控制亮部保护点的y坐标，控制亮部是否过曝,过大可能过曝 0.15
    const double save_point_x = 0.98; // 不建议随便调整
    const double save_point_y = 0.985; // 不建议随便调整
    
    cv::Mat img_gray;
    cv::cvtColor(img_float32, img_gray, cv::COLOR_BGR2GRAY);

    cv::Mat img_norm;
    cv::normalize(img_gray, img_norm, 0, 1, cv::NORM_MINMAX, CV_32F);

    int count_non_zero = cv::countNonZero(star_mask);
    int count_zero = img_norm.total() - count_non_zero;

    cv::Mat img_star_non_zero(cv::Size(1, count_non_zero), CV_32F);
    float* img_star_non_zero_ptr = img_star_non_zero.ptr<float>(0);

    cv::Mat img_bg_non_zero(cv::Size(1, count_zero), CV_32F);
    float* img_bg_non_zero_ptr = img_bg_non_zero.ptr<float>(0);

    std::mutex non_zero_mutex;

    int idx_star = 0;
    int idx_bg = 0;
    star_mask.forEach<float>([&](float& pixel, const int position[]) {
        int y = position[0];
        int x = position[1];

        float* img_norm_ptr = img_norm.ptr<float>(y);

        std::lock_guard<std::mutex> lock(non_zero_mutex);
        if (pixel != 0) {
            img_star_non_zero_ptr[idx_star] = img_norm_ptr[x];
            ++idx_star;
        } else {
            img_bg_non_zero_ptr[idx_bg] = img_norm_ptr[x];
            ++idx_bg;
        }
    });

    double img_star_median = median(img_star_non_zero, 65536, 1.0 + FLT_EPSILON) / 65536.0;
    cv::Scalar img_star_std;
    cv::meanStdDev(img_star_non_zero, cv::noArray(), img_star_std);

    double img_bg_median = median(img_bg_non_zero, 65536, 1.0 + FLT_EPSILON) / 65536.0;
    cv::Scalar img_bg_std;
    cv::meanStdDev(img_bg_non_zero, cv::noArray(), img_bg_std);

    double input_p1 = img_bg_median + point1_LR * img_bg_std.val[0];  // 换作均值
    double output_p1 = input_p1 * point1_UD;

    double input_p2 = img_star_median + point2_LR * img_star_std.val[0];  // 换作均值
    double output_p2 = (input_p2 + point2_UP) <= 0.95 ? (input_p2 + point2_UP) : 0.95;

    std::vector<double> PCHIP_x_point = {0, input_p1, input_p2, save_point_x, 1};
    std::vector<double> PCHIP_y_point = {0, output_p1, output_p2, save_point_y, 1};

    auto spline = boost::math::interpolators::pchip(std::move(PCHIP_x_point), std::move(PCHIP_y_point));

    cv::Mat mv[3], img_transfer_channel[3], img_transfer_channel_flatten[3];
	split(img_BGR_norm, mv);

	for(int k = 0; k < 3; k++) {
        cv::Mat img_flatten = mv[k].reshape(1, 1);

        img_transfer_channel_flatten[k] = cv::Mat(img_flatten.size(), CV_32F);
        float* img_transfer_ptr = img_transfer_channel_flatten[k].ptr<float>(0);

        img_flatten.forEach<float>([&](float& pixel, const int position[]) {
            int x = position[1];
            img_transfer_ptr[x] = spline(pixel);
		});

        img_transfer_channel[k] = img_transfer_channel_flatten[k].reshape(1, mv[k].rows);
    }

	merge(img_transfer_channel, 3, img_transfer);

    return 0;
}

int sKindStretch(const cv::Mat& img_uint16, cv::Mat& img_result) {
    cv::Mat img_float32;
    img_uint16.convertTo(img_float32, CV_32FC3);

    cv::Mat img_dmax_norm;
    cv::normalize(img_float32, img_dmax_norm, 0, 1, cv::NORM_MINMAX, CV_32F);

    cv::Mat img_binary;
    imgSegSStretch(img_dmax_norm, img_binary);

    cv::Mat img_after_S_transform;
    imgSCurveHistogramTransformUseV(img_float32, img_dmax_norm, img_binary, img_after_S_transform);

    img_after_S_transform.convertTo(img_result, CV_16UC3, 65535.0);

    return 0;
}


